'''
Created on 17-Mar-2011

@author: xie
'''
######################## matrix  ##################
import math
import re

import global_var

class SparseVector:
    '''
    A class that implement the sparse vector,
    in which the data is organized as a dic.
    And the dimension to be infinite.
    '''
    vec = None
    
    def __init__(self, **keywords): # dic or list
        for key, value in keywords.iteritems():
            if key == 'dic':
                self.vec = dict(value)
                return
            if key == 'list':
                dic = {}
                i = 0
                for x in value:
                    i = i + 1
                    dic[i] = x
                self.vec = dic
                return
        
        self.vec = {}

        
    def copy(self):
        return SparseVector(dic = self.vec)
    
    def set(self, index, value):
        if value == 0.0:
            if index in self.vec:
                del self.vec[index]
            return
        self.vec[index] = value
        
    def get(self, index):
        if index in self.vec:
            return self.vec[index]
        return 0.0
    
    def inner_product(self, vector):
        s = 0.0
        if self.non_zero_num() <= vector.non_zero_num():
            for id, x in self.vec.iteritems():
                s = s + x*vector.get(id)
        else:
            for id, x in vector.vec.iteritems():
                s = s + x*self.get(id)
                
        return s
    
    def sum(self):
        s = 0
        for v in self.vec.itervalues():
            s = s + v
        return s
    
    def add(self, vector):
        ret = SparseVector()
        s = self.index_set() | vector.index_set()
        for index in s:
            t = self.get(index) + vector.get(index)
            ret.set(index, t)
        return ret
    
    def minus(self, vector):
        ret = SparseVector()
        s = self.index_set() | vector.index_set()
        for index in s:
            t = self.get(index) - vector.get(index)
            ret.set(index, t)
        return ret
    
    def scale(self, value): 
        ret = SparseVector()
        for index, v in self.vec.iteritems():
            ret.set(index, value*v)
        return ret
    
    def index_set(self):
        s = set()
        for index in self.vec.iterkeys():
            s.add(index)
        return s
        
    
    def non_zero_num(self):
        return len(self.vec)
    
    def mod(self):
        s = 0.0
        for x in self.vec.itervalues():
            s = s + x*x
        return math.sqrt(s)
    
    def printInfo(self):
        
        if global_var.logger == None:
            print(self.vec)
        else:
            print >> global_var.logger, self.vec

class SparseMatrix:
    '''
    A class that implement the sparse matrix,
    in which the data is organized as a dic of SparseVector
    And the dimension to be infinite.
    '''
    mat = None
    
    def __init__(self):
        self.mat = {}
        
    def clear(self):
        self.mat.clear()
        
    def copy(self):
        ret = SparseMatrix()
        for id, vec in self.mat.iteritems():
            ret.setRow(id, vec)
        return ret
        
    def transpose(self):
        ret = SparseMatrix()
        for i, row in self.mat.iteritems():
            for j, v in row.vec.iteritems():
                ret.set(j,i,v)
        return ret
    
    def rowOne(self):
        ret = SparseMatrix()
        for i, row in self.mat.iteritems():
            s = row.sum()
            if s != 0: 
                ret.setRow(i, row.scale(1.0/s))
            else:
                ret.setRow(i, row)
        return ret
    
    def modifySuperNode(self, p = 1 - 1e-5): # p: the possibility to transpose to itself
        sn = -1
        for id, vec in self.mat.iteritems():
            t = vec.get(id)
            if t != 0: # super node
                sn = id
                break
        vec = self.getRow(sn)
        s = vec.sum() - 1
        if s == 0:
            return
        vec.set(sn, s*p/(1-p))
        self.setRow(sn, vec)
                
        
    def loadFromMatFile(self, filename):
        self.mat.clear()
        f = file(filename, 'r')
        i = 0
        for line in f:
            i = i + 1
            res = line.split('  ')
            j = 0
            for s in res:
                if s == '':
                    continue
                j = j + 1
                self.set(i,j,float(s))
        f.close()
        
    def loadFromNwbFile(self, filename):
        self.mat.clear()
        f = file(filename, 'r')
        p = re.compile('[0-9]+ [0-9]+ [0-9]+')
        for line in f:
            m = p.match(line)
            if m != None:
                res = line.split(' ')
                i = int(res[0])
                j = int(res[1])
                cnt = int(res[2])
                self.set(i, j, cnt)
        f.close()

    def set(self, i, j, value):
        row = None
        if i in self.mat:
            row = self.mat[i]           
        else:
            row = SparseVector()
            self.mat[i] = row
            
        row.set(j, value)
        
    def setRow(self, i, row):
        vec = row.copy()
        self.mat[i] = vec
    
    def get(self, i, j):
        if i in self.mat:
            row = self.mat[i]
            return row.get(j)
        
        return 0.0
    
    def getRow(self, i):
        vec = self.mat[i]
        return vec.copy()
    
    def map(self, vector):
        ret = SparseVector()
        for i, row in self.mat.iteritems():
            s = row.inner_product(vector)
            ret.set(i, s)
        return ret
    
    def printInfo(self):
        for i, row in self.mat.iteritems():
            
            if global_var.logger == None:
                print("row:" + str(i))
            else:
                print >> global_var.logger, ("row:" + str(i))
            row.printInfo()
            
    def prepareForRW(self):
        ret = self.copy()

        ret.modifySuperNode()       
        ret = ret.rowOne()
        ret = ret.transpose()
        
        return ret
    
    def randomWalkPair(self, c, start_index, deta, file_name = None):
        vec = self.randomWalk(c, start_index, deta)
        
        vec_t = None
        for id in self.mat.iterkeys():
            if id != start_index:
                vec_t = self.randomWalk(c, id, deta)
            else:
                vec_t = vec
                
            v = vec.get(id)
            v *= vec_t.get(start_index)
            vec.set(id, v)
            
        #vec.set(start_index, 0)
        list = sorted(vec.vec.items(), key = lambda d:d[1], reverse = True)
        print_cnt = 30  ##########################
        if file_name != None:
            f = file(file_name, 'w')
            for t in list:
                print >> f, t[0], t[1]
            f.close()
        short_list = list[0:print_cnt]
        
        
        if global_var.logger == None:
            print(short_list)
        else:
            print >> global_var.logger, short_list
        
        return short_list      
            
        
                        
    def randomWalk(self, c, start_index, deta, show = False, file_name = None):
        init_vec = SparseVector(dic = {start_index:1})
        d = 99999
        vec0 = init_vec.copy()
               
        while d > deta :
            vec1 = self.map(vec0)
            vec1 = vec1.scale(c)
            vec1 = vec1.add(init_vec.scale(1-c))
            deta_vec = vec1.minus(vec0)
            d = deta_vec.mod()
            vec0 = vec1
            
        if not show and file_name == None:
            return vec0
            
        list = sorted(vec0.vec.items(), key = lambda d:d[1], reverse = True)
        print_cnt = 30
        if file_name != None:
            f = file(file_name, 'w')
            for t in list:
                f.write(str(t))
            f.close()
        short_list = list[0:print_cnt]
        if show:
            
            if global_var.logger == None:
                print(short_list)
            else:
                print >> global_var.logger, short_list
            
        return vec0
        

 
if __name__ == '__main__':
    '''
    vec1 = SparseVector(list = [0.8147,    0.0975,    0.1576,    0.1419,    0.6557])
    vec2 = SparseVector(list = [0.9058,    0.2785,    0.9706,    0.4218,    0.0357])
    vec3 = SparseVector(list = [0.1270,    0.5469,    0.9572,    0.9157,    0.8491])
    vec4 = SparseVector(list = [0.9134,    0.9575,    0.4854,    0.7922,    0.9340])
    vec5 = SparseVector(list = [0.6324,    0.9649,    0.8003,    0.9595,    0.6787])
    
    vec1 = vec1.scale(0.2)
    vec2 = vec2.scale(0.2)
    vec3 = vec3.scale(0.2)
    vec4 = vec4.scale(0.2)
    vec5 = vec5.scale(0.2)
    
    vec = SparseVector(list = [0.7577,    0.7431,    0.3922,    0.6555,    0.1712])
    
    mat = SparseMatrix()
    mat.setRow(1, vec1)
    mat.setRow(2, vec2)
    mat.setRow(3, vec3)
    mat.setRow(4, vec4)
    mat.setRow(5, vec5)
    '''
    
    mat = SparseMatrix()
    mat.loadFromNwbFile('/mnt/share/graph_2_complete_huixue.nwb')   
    mat = mat.prepareForRW()
    
    mat.randomWalk(0.8, 12, 1e-5, True)
    mat.randomWalkPair(0.8, 12, 1e-5)
    

    